import networkx as nx
from rtree import Rtree
from osgeo import ogr
from mbr import MBR
from shapely.geometry import Point, LineString
from spatial_func import SPoint, distance
import copy
import geopandas as gpd



class UndirRoadNetwork(nx.Graph):
    def __init__(self, g, edge_spatial_idx, edge_idx):
        super(UndirRoadNetwork, self).__init__(g)
        # entry: eid
        self.edge_spatial_idx = edge_spatial_idx
        # eid -> edge key (start_coord, end_coord)
        self.edge_idx = edge_idx

    def to_directed(self, as_view=False):
        """
        Convert undirected road network to directed road network
        new edge will have new eid, and each original edge will have two edge with reversed coords
        :return:
        """
        assert as_view is False, "as_view is not supported"
        avail_eid = max([eid for u, v, eid in self.edges.data(data='eid')]) + 1
        g = nx.DiGraph()
        edge_spatial_idx = Rtree()
        edge_idx = {}
        # add nodes
        for n, data in self.nodes(data=True):
            # when data=True, it means will data=node's attributes
            new_data = copy.deepcopy(data)
            g.add_node(n, **new_data)
        # add edges
        for u, v, data in self.edges(data=True):
            mbr = MBR.cal_mbr(data['coords'])
            # add forward edge
            forward_data = copy.deepcopy(data)
            g.add_edge(u, v, **forward_data)
            edge_spatial_idx.insert(forward_data['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
            edge_idx[forward_data['eid']] = (u, v)
            # add backward edge
            backward_data = copy.deepcopy(data)
            backward_data['eid'] = avail_eid
            avail_eid += 1
            backward_data['coords'].reverse()
            g.add_edge(v, u, **backward_data)
            edge_spatial_idx.insert(backward_data['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
            edge_idx[backward_data['eid']] = (v, u)
        print('# of nodes:{}'.format(g.number_of_nodes()))
        print('# of edges:{}'.format(g.number_of_edges()))
        return RoadNetwork(g, edge_spatial_idx, edge_idx)

    def range_query(self, mbr):
        """
        spatial range query. Given a mbr, return a range of edges.
        :param mbr: query mbr
        :return: qualified edge keys
        """
        eids = self.edge_spatial_idx.intersection((mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        return [self.edge_idx[eid] for eid in eids]

    def remove_edge(self, u, v):
        edge_data = self[u][v]
        coords = edge_data['coords']
        mbr = MBR.cal_mbr(coords)
        # delete self.edge_idx[eid] from edge index
        del self.edge_idx[edge_data['eid']]
        # delete from spatial index
        self.edge_spatial_idx.delete(edge_data['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        # delete from graph
        super(UndirRoadNetwork, self).remove_edge(u, v)

    def add_edge(self, u_of_edge, v_of_edge, **attr):
        coords = attr['coords']
        mbr = MBR.cal_mbr(coords)
        attr['length'] = sum([distance(coords[i], coords[i + 1]) for i in range(len(coords) - 1)])
        # add edge to edge index
        self.edge_idx[attr['eid']] = (u_of_edge, v_of_edge)
        # add edge to spatial index
        self.edge_spatial_idx.insert(attr['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        # add edge to graph
        super(UndirRoadNetwork, self).add_edge(u_of_edge, v_of_edge, **attr)


class RoadNetwork(nx.DiGraph):
    def __init__(self, g, edge_spatial_idx, edge_idx):
        super(RoadNetwork, self).__init__(g)
        # entry: eid
        self.edge_spatial_idx = edge_spatial_idx
        # eid -> edge key (start_coord, end_coord)
        self.edge_idx = edge_idx

    def range_query(self, mbr):
        """
        spatial range query
        :param mbr: query mbr
        :return: qualified edge keys
        """
        eids = self.edge_spatial_idx.intersection((mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        return [self.edge_idx[eid] for eid in eids]

    def remove_edge(self, u, v):
        edge_data = self[u][v]
        coords = edge_data['coords']
        mbr = MBR.cal_mbr(coords)
        # delete self.edge_idx[eifrom edge index
        del self.edge_idx[edge_data['eid']]
        # delete from spatial index
        self.edge_spatial_idx.delete(edge_data['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        # delete from graph
        super(RoadNetwork, self).remove_edge(u, v)

    def add_edge(self, u_of_edge, v_of_edge, **attr):
        coords = attr['coords']
        mbr = MBR.cal_mbr(coords)
        attr['length'] = sum([distance(coords[i], coords[i + 1]) for i in range(len(coords) - 1)])
        # add edge to edge index
        self.edge_idx[attr['eid']] = (u_of_edge, v_of_edge)
        # add edge to spatial index
        self.edge_spatial_idx.insert(attr['eid'], (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        # add edge to graph
        super(RoadNetwork, self).add_edge(u_of_edge, v_of_edge, **attr)


def load_rn_shp(path, is_directed=True):
    edge_spatial_idx = Rtree()
    edge_idx = {}
    # 使用 geopandas 读取 Shapefile
    gdf = gpd.read_file(path)
    # 确保坐标系为 WGS84 (EPSG:4326)
    if gdf.crs.to_epsg() != 4326:
        gdf = gdf.to_crs(epsg=4326)
    # 创建图
    g = nx.DiGraph() if is_directed else nx.Graph()
    # 处理每条边
    for idx, row in gdf.iterrows():
        geom = row.geometry
        # 修正坐标访问方式，并转换为 (纬度, 经度) 元组
        coords = [(pt[1], pt[0]) for pt in geom.coords]  # (lat, lng)
        start_node = coords[0]
        end_node = coords[-1]
        # 添加节点（SPoint 需恢复为 (经度, 纬度)）
        g.add_node(start_node, pt=SPoint(start_node[1], start_node[0]))
        g.add_node(end_node, pt=SPoint(end_node[1], end_node[0]))
        # 添加边属性
        eid = row['osm_id']  # 根据实际字段名调整
        edge_attrs = {
            'eid': eid,
            'coords': [SPoint(pt[1], pt[0]) for pt in coords],  # 转换为 SPoint(lng, lat)
            'length': sum([distance(SPoint(coords[i][1], coords[i][0]),
                          SPoint(coords[i+1][1], coords[i+1][0]))
                        for i in range(len(coords)-1)])
        }
        g.add_edge(start_node, end_node, **edge_attrs)
        # 更新空间索引
        mbr = MBR.cal_mbr(edge_attrs['coords'])
        edge_spatial_idx.insert(eid, (mbr.min_lng, mbr.min_lat, mbr.max_lng, mbr.max_lat))
        edge_idx[eid] = (start_node, end_node)
    print('# of nodes:{}'.format(g.number_of_nodes()))
    print('# of edges:{}'.format(g.number_of_edges()))
    if is_directed:
        return RoadNetwork(g, edge_spatial_idx, edge_idx)
    else:
        return UndirRoadNetwork(g, edge_spatial_idx, edge_idx)


def store_rn_shp(rn, target_path):
    print('# of nodes:{}'.format(rn.number_of_nodes()))
    print('# of edges:{}'.format(rn.number_of_edges()))

    # ----------- 保存节点 -----------
    # 提取节点坐标和属性
    nodes = []
    for node_id, data in rn.nodes(data=True):
        # 假设节点坐标以元组 (lng, lat) 形式存储
        geometry = Point(node_id[0], node_id[1])  # 或根据实际数据结构调整
        nodes.append({
            'nid': node_id,  # 节点唯一标识符
            'geometry': geometry,
            # 可添加其他属性（如 data 中的字段）
        })

    # 创建节点 GeoDataFrame
    gdf_nodes = gpd.GeoDataFrame(nodes, crs='EPSG:4326')  # 假设坐标系为 WGS84

    # ----------- 保存边 -----------
    # 提取边几何和属性
    edges = []
    for u, v, data in rn.edges(data=True):
        # 从边数据中提取坐标（假设 data['coords'] 是 SPoint 列表）
        line_coords = [(sp.lng, sp.lat) for sp in data.get('coords', [])]
        geometry = LineString(line_coords)
        edges.append({
            'eid': data.get('eid', 'N/A'),  # 边唯一标识符
            'length': data.get('length', 0),
            'geometry': geometry,
            # 可添加其他属性（如 u, v 等）
        })

    # 创建边 GeoDataFrame
    gdf_edges = gpd.GeoDataFrame(edges, crs='EPSG:4326')

    # ----------- 保存为 Shapefile -----------
    # 节点和边分开保存（兼容 nx.write_shp 行为）
    gdf_nodes.to_file(f"{target_path}nodes.shp")
    gdf_edges.to_file(f"{target_path}edges.shp")


if __name__ == '__main__':
    undir_rn = load_rn_shp(r"D:\deposite\Diff-RNTraj\data\Tokyo\road_shp\roads.shp", is_directed=False)
    mbr = MBR(min_lng=139.47, min_lat=35.58, max_lng=139.80, max_lat=35.80)
    dir_rn = undir_rn.to_directed()
    store_rn_shp(dir_rn, r"D:\deposite\Diff-RNTraj\data\Tokyo\road_net")